import os
from glob import glob
import torch
import matplotlib.pyplot as plt
import numpy as np

from nets.PrismaMattingNet import PrismaNet
from loss import HairMatLoss, dice_loss, CrossEntropyLoss2d, iou_loss, fusion_loss
import utils

class Trainer:
    def __init__(self, config, dataloader):
        self.config = config
        self.lr = config.lr
        self.epoch = config.epoch
        self.checkpoint_dir = config.checkpoint_dir
        self.model_path = config.checkpoint_dir
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.data_loader = dataloader
        self.image_len = len(dataloader)
        self.num_classes = config.num_classes
        self.build_model()
        self.sample_step = config.sample_step
        self.sample_dir = config.sample_dir

    def build_model(self):
        self.net = PrismaNet()
        self.net.to(self.device)
        if self.config.snapshot!='':
            self.net.load_state_dict(torch.load(model[-1], map_location=self.device))


    def train(self):
        loss_value = utils.AverageMeter()
        iou_value = utils.AverageMeter()

        MobileHairNetLoss = HairMatLoss().to(self.device)
        optimizer = torch.optim.Adam(self.net.parameters(), lr=self.lr, eps=1e-7)

        for epoch in range(self.max_epoch):
            loss_value.reset()
            for step, (image, mask, _, _) in enumerate(self.train_data_loader):
                image = image.to(self.device)
                mask = mask.to(self.device)
                pred = self.net(image)
                self.net.zero_grad()
                # loss = dice_loss(pred, mask)
                loss = fusion_loss(pred, mask, image)
                loss.backward()
                optimizer.step()

                iou = iou_loss(pred, mask)

                loss_value.update(loss)
                iou_value.update(iou)

                print("epoch: [%d/%d] | image: [%d/%d] | loss: %.4f | IOU: %.4f" % (epoch, self.max_epoch, step, self.image_len, loss_value.avg, iou_value.avg))

                # save sample images
                if step % self.sample_step == 0:
                    # self.save_sample_imgs(image[0], mask[0], torch.argmax(pred[0], 0), self.sample_dir, epoch, step)
                    utils.save_sample_imgs(image[0], mask[0], pred[0], self.sample_dir, epoch, step)
                    print('[*] Saved sample images')

            torch.save(self.net.state_dict(), '%s/PrismaMattingNet_portrait_epoch-%04d.pth' % (self.checkpoint_dir, epoch))
